#include "Init.h"
#include "Rules.h"
#include "Utils.h"

ERESOURCE g_RulesLock = {0};
UNICODE_STRING g_usRegistryPath = {0};

PSINGLE_LIST_ENTRY g_pRulesListHead = NULL;

// actions definitions
UNICODE_STRING g_usHideFolderActionValue = {0}; // HideFolderAction

NTSTATUS InitializeRulesManager(PUNICODE_STRING pusRegistryPath)
{
    NTSTATUS ntStatus = ExInitializeResourceLite(&g_RulesLock);

    static UNICODE_STRING usRulesKey = {0};

    RtlInitUnicodeString(&usRulesKey, L"\\Rules");

    ntStatus = ConcatUnicodeString(pusRegistryPath, &usRulesKey, &g_usRegistryPath, PagedPool);

    if (!NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    }

    // initialize actions
    RtlInitUnicodeString(&g_usHideFolderActionValue, L"HideFolderAction");

    return ntStatus;
}

NTSTATUS FreeRulesList()
{
    NTSTATUS ntStatus = AcquireResource(&g_RulesLock, TRUE);

    if (!NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    }

    while (NULL != g_pRulesListHead)
    {
        PSINGLE_LIST_ENTRY pRemovedEntry = PopEntryList(g_pRulesListHead);
        PRULE pTempRule = NULL;

        if (NULL == pRemovedEntry) // list is empty and has only one record
        {
            pRemovedEntry = g_pRulesListHead;
            g_pRulesListHead = NULL;
        }
        
        pTempRule = CONTAINING_RECORD(pRemovedEntry, RULE, singleListEntry);

        switch (pTempRule->ruleCategory)
        {
            case HideFolderAction:
            {
                ExFreePoolWithTag(pTempRule->HideFolderRule.usFolderForHide.Buffer, FileHidingFilterPoolTag);
                ExFreePoolWithTag(pTempRule->HideFolderRule.usParentFolder.Buffer, FileHidingFilterPoolTag);
                ExFreePoolWithTag(pTempRule, FileHidingFilterPoolTag);

                break;
            }
            default:
            {
                ASSERT(FALSE);
            }
        }
    }

    ReleaseResource(&g_RulesLock);

    return STATUS_SUCCESS;
}

NTSTATUS RuleBuilderCallback(PUNICODE_STRING pusValue, const wchar_t* pwsPath);

NTSTATUS ReloadRules()
{
    NTSTATUS ntStatus = AcquireResource(&g_RulesLock, TRUE);

    if (!NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    }

    if (NULL != g_pRulesListHead)
    {
        ntStatus = FreeRulesList();

        ASSERT(NT_SUCCESS(ntStatus));
    }

    // load HideFolderAction set
    {
        ntStatus = EnumerateMultiSZValue(&g_usRegistryPath, &g_usHideFolderActionValue, RuleBuilderCallback);

        if (!NT_SUCCESS(ntStatus))
        {
            ReleaseResource(&g_RulesLock);

            return ntStatus;
        }
    }

    ReleaseResource(&g_RulesLock);

    return STATUS_SUCCESS;
}

NTSTATUS RuleBuilderCallback(PUNICODE_STRING pusValue, const wchar_t* pwsPath)
{
    const wchar_t wsSymLinkPrefix[4] = {L'\\', L'?', L'?', L'\\'};
    HANDLE hDevice = NULL;
    NTSTATUS ntStatus = STATUS_SUCCESS;
    OBJECT_ATTRIBUTES objAttributes = {0};
    IO_STATUS_BLOCK ioStatus = {0};
    PFILE_OBJECT pFileObject = NULL;
    PDEVICE_OBJECT pRelatedDeviceObject = NULL;
    PDEVICE_OBJECT pLowestDeviceObject = NULL;
    PRULE pRule = NULL;
    UNICODE_STRING usKernelDisk = {0};
    UNICODE_STRING usParentFolder = {0};
    UNICODE_STRING usFolderForHide = {0};
    UNICODE_STRING usHide = {0};
    wchar_t* pwsDevice = 0;
    wchar_t* pwsRootFolder = 0;
    wchar_t* pwsHiddenFolder = 0;

    // parse path for Device

    const wchar_t* pwsDeviceEnd = wcschr(pwsPath, L'\\'); 
    const wchar_t* pwsRootFolderBegin = wcschr(pwsPath, L'\\');
    const wchar_t* pwsRootFolderEnd = wcsrchr(pwsPath, L'\\');
    const wchar_t* pwsHiddenFolderBegin = wcsrchr(pwsPath, L'\\') + 1;

    size_t uiDeviceStringLen = 0;
    size_t uiRootFolderStringLen = 0;
    size_t uiHiddenFolderSize = 0;

    if (0 == pwsDeviceEnd)
    {
        return STATUS_INVALID_PARAMETER_2;
    }

    if (0 == pwsRootFolderBegin)
    {
        return STATUS_INVALID_PARAMETER_2;
    }

    if (L'\0' == *pwsHiddenFolderBegin)
    {
        return STATUS_INVALID_PARAMETER_2;
    }

    uiDeviceStringLen = (char*)pwsDeviceEnd - (char*)pwsPath + sizeof(wchar_t) /*last "\\"*/ + sizeof(wsSymLinkPrefix) /*\\??\\*/ + sizeof(wchar_t) /*\0*/;
    uiRootFolderStringLen = (char*)pwsRootFolderEnd - (char*)pwsRootFolderBegin;
    uiHiddenFolderSize = wcslen(pwsHiddenFolderBegin) * sizeof(wchar_t) + sizeof(wchar_t) /*\0*/;

    //Getting device object from disk name

    pwsDevice = (wchar_t*)ExAllocatePoolWithTag(NonPagedPool, uiDeviceStringLen, FileHidingFilterPoolTag);

    memset((char*)pwsDevice, 0, uiDeviceStringLen);

    memcpy((char*)pwsDevice, &wsSymLinkPrefix[0], sizeof(wsSymLinkPrefix));
    memcpy((char*)pwsDevice + sizeof(wsSymLinkPrefix), (char*)pwsPath, (char*)pwsDeviceEnd - (char*)pwsPath + sizeof(wchar_t));

    RtlInitUnicodeString(&usKernelDisk, pwsDevice);

    //Getting parent folder name

    if (0 == uiRootFolderStringLen) /*only one "\\" found, example c:\\Secret*/
    {
        uiRootFolderStringLen = sizeof(wchar_t);/*\\*/
    }

    uiRootFolderStringLen += sizeof(wchar_t); /*\0*/

    pwsRootFolder = (wchar_t*)ExAllocatePoolWithTag(NonPagedPool,uiRootFolderStringLen, FileHidingFilterPoolTag);

    memset((char*)pwsRootFolder, 0, uiRootFolderStringLen);
    memcpy((char*)pwsRootFolder, (char*)pwsRootFolderBegin, uiRootFolderStringLen - sizeof(wchar_t)); /*exclude last \0*/

    RtlInitUnicodeString(&usParentFolder, pwsRootFolder);

    //Getting name of folder for hide
    pwsHiddenFolder = (wchar_t*)ExAllocatePoolWithTag(NonPagedPool, uiHiddenFolderSize, FileHidingFilterPoolTag); 

    memset(pwsHiddenFolder, 0, uiHiddenFolderSize);
    memcpy((char*)pwsHiddenFolder, (char*)pwsHiddenFolderBegin, uiHiddenFolderSize /*- sizeof(wchar_t) exclude \0*/);

    RtlInitUnicodeString(&usFolderForHide, pwsHiddenFolder);   

    //Getting the lowest device object

    InitializeObjectAttributes(&objAttributes, &usKernelDisk, OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, NULL, NULL );

    ntStatus = ZwCreateFile(&hDevice,
                            SYNCHRONIZE,
                            &objAttributes,
                            &ioStatus,
                            NULL,
                            0,
                            FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
                            FILE_OPEN,
                            FILE_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT,
                            NULL,
                            0);

    if (!NT_SUCCESS(ntStatus))
    {
        if(NULL != pwsRootFolder)
        {
            ExFreePoolWithTag(pwsRootFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsHiddenFolder)
        {
            ExFreePoolWithTag(pwsHiddenFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsDevice)
        {
            ExFreePoolWithTag(pwsDevice, FileHidingFilterPoolTag);
        }

        ZwClose(hDevice);

        return ntStatus;
    }

    ExFreePoolWithTag(pwsDevice, FileHidingFilterPoolTag);

    ntStatus = ObReferenceObjectByHandle(hDevice,
                                         0,
                                         *IoFileObjectType,
                                         KernelMode,
                                         (PVOID*)&pFileObject,
                                         NULL);

    ZwClose(hDevice);

    if (!NT_SUCCESS(ntStatus))
    {

        if(NULL != pwsRootFolder)
        {
            ExFreePoolWithTag(pwsRootFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsHiddenFolder)
        {
            ExFreePoolWithTag(pwsHiddenFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsDevice)
        {
            ExFreePoolWithTag(pwsDevice, FileHidingFilterPoolTag);
        }



        return ntStatus;
    }

    pRelatedDeviceObject = IoGetRelatedDeviceObject(pFileObject);
    pLowestDeviceObject = IoGetDeviceAttachmentBaseRef(pRelatedDeviceObject);

    if (pRelatedDeviceObject == pLowestDeviceObject)
    {
        ASSERT(FALSE);

        if(NULL != pwsRootFolder)
        {
            ExFreePoolWithTag(pwsRootFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsHiddenFolder)
        {
            ExFreePoolWithTag(pwsHiddenFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsDevice)
        {
            ExFreePoolWithTag(pwsDevice, FileHidingFilterPoolTag);
        }


        // maybe ObDereferenceObject???

        return STATUS_INVALID_PARAMETER_2;
    }

    ObDereferenceObject(pFileObject);
    ObDereferenceObject(pLowestDeviceObject);

    //Creation of new rule

    pRule =(PRULE)ExAllocatePoolWithTag(NonPagedPool, sizeof(RULE), FileHidingFilterPoolTag);

    if (0 == pRule)
    {
        if(NULL != pwsRootFolder)
        {
            ExFreePoolWithTag(pwsRootFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsHiddenFolder)
        {
            ExFreePoolWithTag(pwsHiddenFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsDevice)
        {
            ExFreePoolWithTag(pwsDevice, FileHidingFilterPoolTag);
        }

        return STATUS_INSUFFICIENT_RESOURCES;
    }

    memset(pRule, 0, sizeof(RULE));

    RtlInitUnicodeString(&usHide, L"HideFolderAction");

    if(RtlEqualUnicodeString(pusValue, &usHide, TRUE))
    {
        pRule->ruleCategory = HideFolderAction;
        pRule->HideFolderRule.pDeviceObject = pLowestDeviceObject;
        pRule->HideFolderRule.usParentFolder = usParentFolder;
        pRule->HideFolderRule.usFolderForHide = usFolderForHide;
    }

    else
    {
        if(NULL != pwsRootFolder)
        {
            ExFreePoolWithTag(pwsRootFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsHiddenFolder)
        {
            ExFreePoolWithTag(pwsHiddenFolder, FileHidingFilterPoolTag);
        }

        if(NULL != pwsDevice)
        {
            ExFreePoolWithTag(pwsDevice, FileHidingFilterPoolTag);
        }

        if(NULL != pRule)
        {
            ExFreePoolWithTag(pRule, FileHidingFilterPoolTag);
        }

        return STATUS_INVALID_PARAMETER_1;
    }

    if (NULL == g_pRulesListHead)
    {
        // this is first loaded rule
        g_pRulesListHead = &pRule->singleListEntry;
    }
    else
    {
        PushEntryList(g_pRulesListHead, &pRule->singleListEntry);
    }

    return ntStatus;
}

NTSTATUS ProcessRules(IN PDEVICE_OBJECT pDeviceObject, 
                      IN PUNICODE_STRING pusParentDirectory,
                      IN PUNICODE_STRING pusFolderForHide,
                      OUT PRULE pRuleFound)
{
    PSINGLE_LIST_ENTRY pListEntryIterator = NULL;
    PRULE pTempRule = NULL;
    NTSTATUS ntStatus = AcquireResource(&g_RulesLock, FALSE);

    if (!NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    }

    RtlZeroMemory(pRuleFound, sizeof(RULE));

    pRuleFound->ruleCategory = AllowAction;

    if (NULL == g_pRulesListHead)
    {
        pRuleFound->ruleCategory = AllowAction;

        ReleaseResource(&g_RulesLock);

        return STATUS_SUCCESS;
    }

    pListEntryIterator = g_pRulesListHead;

    for (; NULL != pListEntryIterator; pListEntryIterator = pListEntryIterator->Next)
    {
        pTempRule = CONTAINING_RECORD(pListEntryIterator, RULE, singleListEntry);

        switch (pTempRule->ruleCategory)
        {
            case HideFolderAction:
            {
                if(pDeviceObject != pTempRule->HideFolderRule.pDeviceObject)
                {
                    continue;
                }

                if(RtlCompareUnicodeString(&(pTempRule->HideFolderRule.usParentFolder), pusParentDirectory, TRUE) != 0)
                {
                    continue;
                }

                if(RtlCompareUnicodeString(&(pTempRule->HideFolderRule.usFolderForHide), pusFolderForHide, TRUE) != 0)
                {
                    continue;
                }

                *pRuleFound = *pTempRule;

                ReleaseResource(&g_RulesLock);

                return STATUS_SUCCESS;
            }
        }
    }

    ReleaseResource(&g_RulesLock);

    return STATUS_SUCCESS;
}